import os
import subprocess
import argparse
import re

def detect_gpu_type():
    """检测系统可用的GPU类型"""
    try:
        subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return 'nvidia'
    except (subprocess.CalledProcessError, FileNotFoundError):
        try:
            subprocess.run(['rocminfo'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            return 'amd'
        except (subprocess.CalledProcessError, FileNotFoundError):
            try:
                subprocess.run(['intel_gpu_top'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
                return 'intel'
            except (subprocess.CalledProcessError, FileNotFoundError):
                return None

def get_gpu_encoder(gpu_type, codec='h264'):
    """根据GPU类型返回最佳编码器"""
    if gpu_type == 'nvidia':
        return 'h264_nvenc' if codec == 'h264' else 'hevc_nvenc'
    elif gpu_type == 'amd':
        return 'h264_amf' if codec == 'h264' else 'hevc_amf'
    elif gpu_type == 'intel':
        return 'h264_qsv' if codec == 'h264' else 'hevc_qsv'
    else:
        return 'libx264' if codec == 'h264' else 'libx265'

def convert_webm_to_mp4(input_file, output_file=None, overwrite=False, gpu_accel=True, 
                         quality='high', codec='h264', audio_codec='aac'):
    """
    将WebM文件转换为MP4格式，使用GPU加速并优化质量
    
    参数:
    input_file: 输入的WebM文件路径
    output_file: 输出的MP4文件路径（默认为输入文件同目录同名）
    overwrite: 是否覆盖已存在的输出文件
    gpu_accel: 是否使用GPU加速
    quality: 视频质量设置（'low', 'medium', 'high'）
    codec: 视频编码格式（'h264'或'hevc'）
    audio_codec: 音频编码格式（'aac'或'copy'）
    """
    # 处理Windows长路径问题
    if os.name == 'nt':
        input_file = r'\\?\\' + os.path.abspath(input_file)
    
    # 设置默认输出文件名
    if output_file is None:
        filename, _ = os.path.splitext(input_file)
        output_file = filename + ".mp4"
    else:
        if os.name == 'nt':
            output_file = r'\\?\\' + os.path.abspath(output_file)
    
    # 检查输入文件是否存在
    if not os.path.isfile(input_file):
        # 尝试使用短路径格式
        try:
            # 获取短路径名
            import ctypes
            buf = ctypes.create_unicode_buffer(1024)
            ctypes.windll.kernel32.GetShortPathNameW(input_file, buf, 1024)
            input_file = buf.value
        except:
            pass
        
        if not os.path.isfile(input_file):
            print(f"错误: 输入文件不存在 - {input_file}")
            print("请检查以下可能原因:")
            print("1. 文件路径是否正确")
            print("2. 文件名是否包含特殊字符")
            print("3. 文件名是否过长（尝试缩短文件名）")
            print("4. 文件是否被其他程序占用")
            return False
    
    # 处理已存在的输出文件
    if os.path.isfile(output_file):
        if overwrite:
            try:
                os.remove(output_file)
            except Exception as e:
                print(f"无法删除已存在的输出文件: {e}")
                return False
        else:
            print(f"错误: 输出文件已存在 - {output_file} (使用 -y 选项覆盖)")
            return False
    
    # 检测GPU类型
    gpu_type = detect_gpu_type() if gpu_accel else None
    encoder = get_gpu_encoder(gpu_type, codec) if gpu_accel else ('libx264' if codec == 'h264' else 'libx265')
    
    print(f"使用视频编码器: {encoder} ({gpu_type or 'CPU'})")
    print(f"视频质量: {quality}")
    print(f"视频编码格式: {codec.upper()}")
    print(f"音频编码格式: {audio_codec}")
    
    # 构建FFmpeg命令
    cmd = [
        'ffmpeg',
        '-i', input_file,
        '-c:v', encoder,
        '-movflags', '+faststart',
        '-y' if overwrite else '-n',
    ]
    
    # 设置音频编码
    if audio_codec == 'copy':
        cmd.extend(['-c:a', 'copy'])
    else:
        cmd.extend(['-c:a', audio_codec])
        if audio_codec == 'aac':
            cmd.extend(['-b:a', '256k'])
    
    # 根据质量设置添加参数
    if quality == 'high':
        if 'nvenc' in encoder:
            cmd.extend(['-preset', 'p7', '-tune', 'hq', '-cq', '18'])
        elif 'amf' in encoder:
            cmd.extend(['-quality', 'quality', '-qp_i', '18', '-qp_p', '18'])
        elif 'qsv' in encoder:
            cmd.extend(['-preset', 'slower', '-global_quality', '18'])
        else:  # CPU编码
            cmd.extend(['-preset', 'slow', '-crf', '18', '-tune', 'film'])
    elif quality == 'medium':
        if 'nvenc' in encoder:
            cmd.extend(['-preset', 'p7', '-cq', '22'])
        elif 'amf' in encoder:
            cmd.extend(['-quality', 'balanced', '-qp_i', '22', '-qp_p', '22'])
        elif 'qsv' in encoder:
            cmd.extend(['-preset', 'medium', '-global_quality', '22'])
        else:  # CPU编码
            cmd.extend(['-preset', 'medium', '-crf', '22', '-tune', 'film'])
    else:  # low
        if 'nvenc' in encoder:
            cmd.extend(['-preset', 'p7', '-cq', '26'])
        elif 'amf' in encoder:
            cmd.extend(['-quality', 'speed', '-qp_i', '26', '-qp_p', '26'])
        elif 'qsv' in encoder:
            cmd.extend(['-preset', 'fast', '-global_quality', '26'])
        else:  # CPU编码
            cmd.extend(['-preset', 'fast', '-crf', '26', '-tune', 'film'])
    
    cmd.append(output_file)
    
    try:
        print("开始转换...")
        process = subprocess.Popen(
            cmd,
            stderr=subprocess.PIPE,
            universal_newlines=True,
            encoding='utf-8',
            errors='replace'
        )
        
        # 实时显示进度
        while True:
            line = process.stderr.readline()
            if not line:
                break
            if 'time=' in line:
                time_match = re.search(r'time=(\d+:\d+:\d+\.\d+)', line)
                if time_match:
                    print(f"\r转换进度: {time_match.group(1)}", end='', flush=True)
        
        process.wait()
        
        if process.returncode == 0:
            print(f"\n转换成功: {input_file} -> {output_file}")
            return True
        else:
            print(f"\n转换失败，返回码: {process.returncode}")
            print("错误信息:")
            while True:
                line = process.stderr.readline()
                if not line:
                    break
                print(line, end='')
            return False
    except FileNotFoundError:
        print("\n错误: 未找到FFmpeg，请先安装FFmpeg并添加到系统路径")
        return False
    except Exception as e:
        print(f"\n转换过程中发生错误: {e}")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='使用GPU加速将WebM文件转换为MP4格式')
    parser.add_argument('input', help='输入WebM文件路径')
    parser.add_argument('-o', '--output', help='输出MP4文件路径（可选）')
    parser.add_argument('-y', '--overwrite', action='store_true', help='覆盖已存在的输出文件')
    parser.add_argument('--no-gpu', action='store_true', help='禁用GPU加速（使用CPU编码）')
    parser.add_argument('-q', '--quality', choices=['low', 'medium', 'high'], default='high',
                        help='视频质量设置（low, medium, high），默认为high')
    parser.add_argument('-c', '--codec', choices=['h264', 'hevc'], default='hevc',
                        help='视频编码格式（h264, hevc），默认为hevc')
    parser.add_argument('-a', '--audio', choices=['aac', 'copy'], default='aac',
                        help='音频处理方式（aac: 转码为AAC, copy: 直接复制），默认为aac')

    args = parser.parse_args()
    
    # 处理Windows路径的特殊字符
    if os.name == 'nt':
        args.input = r'{}'.format(args.input)
    
    convert_webm_to_mp4(
        input_file=args.input,
        output_file=args.output,
        overwrite=args.overwrite,
        gpu_accel=not args.no_gpu,
        quality=args.quality,
        codec=args.codec,
        audio_codec=args.audio
    )